!==========================================================
! This procedure is used to calculate the kinetic spectra 
! using the DCT (Discrete Cosine Transform) method of 
! Denis et al. (2002).
! Author: Jianghao Li
!   2023.11
!==========================================================
program calc_kinetic_spectra
  
  use netcdf

  implicit none
  
  integer, parameter  :: r8     = 8
  real(r8), parameter :: radius = 6.37122d6
  real(r8), parameter :: pi     = atan(1.0_r8) * 4.0_r8
  real(r8), parameter :: rad    = pi / 180.0_r8
  character(len=*), parameter :: file_in  = "../MCV/mcv_uv200hPa_regional4.nc"
  character(len=*), parameter :: file_out = "../uv_regional4_mcv.nc"
  
  character(len=*), parameter :: lat_name = "lat"
  character(len=*), parameter :: lon_name = "lon"
  character(len=*), parameter :: u_name   = "u"
  character(len=*), parameter :: v_name   = "v"
  character(len=*), parameter :: UNITS    = "units"
  
  logical, parameter :: vor_div_based = .false.

  integer :: ncid
  integer :: lat_dimid, lon_dimid
  integer :: nlats_in, nlons_in
  integer :: nlats, nlons

  integer :: lat_varid, lon_varid
  integer :: u_varid, v_varid
  integer :: ndims_in, nvars_in, ngatts_in, unlimdimid_in

  real(r8), allocatable :: lats(:), lons(:)
  real(r8), allocatable :: u(:,:), v(:,:)
  real(r8), allocatable :: div(:,:), vor(:,:)
  real(r8), allocatable :: DCT_div(:,:), DCT_vor(:,:)

  integer rec, start(2), count(2), Nwave
  real(r8) dlon, dlat
  real(r8), allocatable :: power_div(:), power_vor(:), wavelength(:), wavenumber(:)
  integer wavenumber_dimid, div_varid, vor_varid, powerdiv_varid, powervor_varid, wavelength_varid, wavenumber_varid

  call readdata()

  nlons = nlons_in - 2
  nlats = nlats_in - 2
  Nwave = min(nlons, nlats) - 1
  allocate(div(1:nlons, 1:nlats), vor(1:nlons, 1:nlats))
  allocate(DCT_div(0:nlons-1, 0:nlats-1), DCT_vor(0:nlons-1, 0:nlats-1))
  allocate(power_div(Nwave), power_vor(Nwave), wavelength(Nwave), wavenumber(Nwave))

  !-------------main computational procedure ----------------
  if (vor_div_based) then
    call calc_div_vor(u, v, div, vor)
    call calc_DCT(div, vor, DCT_div, DCT_vor)
  else
    call calc_DCT(u(2:nlons_in-1, 2:nlats_in-1), v(2:nlons_in-1, 2:nlats_in-1), DCT_div, DCT_vor)
  end if
  call calc_power_spectral(DCT_div, DCT_vor, power_div, power_vor, wavelength, wavenumber)

  call output()

  deallocate(u, v)
  deallocate(div, vor)
  deallocate(lats, lons)
  deallocate(DCT_div, DCT_vor)
  deallocate(power_div, power_vor, wavelength)

contains
  subroutine readdata()
  
    call check(nf90_open(File_in, nf90_nowrite, ncid))
    call check(nf90_inquire(ncid, ndims_in, nvars_in, ngatts_in, unlimdimid_in))

    call check(nf90_inq_dimid(ncid, lat_name, lat_dimid))
    call check(nf90_inquire_dimension(ncid, lat_dimid, len=nlats_in))
    call check(nf90_inq_dimid(ncid, lon_name, lon_dimid))
    call check(nf90_inquire_dimension(ncid, lon_dimid, len=nlons_in))

    print*, 'nlats_in, nlons_in =', nlats_in, nlons_in

    allocate(lats(nlats_in), lons(nlons_in))
    call check(nf90_inq_varid(ncid, lat_name, lat_varid))
    call check(nf90_inq_varid(ncid, lon_name, lon_varid))
    call check(nf90_get_var(ncid, lat_varid, lats))
    call check(nf90_get_var(ncid, lon_varid, lons))

    dlon = (lons(2) - lons(1)) * rad
    dlat = (lats(2) - lats(1)) * rad
       
    allocate(u(nlons_in, nlats_in), &
             v(nlons_in, nlats_in))
    call check(nf90_inq_varid(ncid, u_name, u_varid))
    call check(nf90_inq_varid(ncid, v_name, v_varid))

    count = (/nlons_in, nlats_in/)
    start = (/1, 1/)
    !do rec = 2, 2 !ntimes
    !  start(4) = rec
    call check(nf90_get_var(ncid, u_varid, u, start=start, count=count))
    call check(nf90_get_var(ncid, v_varid, v, start=start, count=count))
    !end do
    call check(nf90_close(ncid))
 
  end subroutine readdata

  subroutine output()
  
    call check(nf90_create(file_out, nf90_clobber, ncid))
    call check(nf90_def_dim(ncid, "wavenumber", Nwave, wavenumber_dimid))
    call check(nf90_def_dim(ncid, "lat"       , nlats, lat_dimid))
    call check(nf90_def_dim(ncid, "lon"       , nlons, lon_dimid))

    call check(nf90_def_var(ncid, "wavenumber"    , nf90_double, wavenumber_dimid, wavenumber_varid))
    call check(nf90_def_var(ncid, "lat"           , nf90_double, lat_dimid, lat_varid))
    call check(nf90_def_var(ncid, "lon"           , nf90_double, lon_dimid, lon_varid))
    call check(nf90_put_att(ncid, wavenumber_varid, UNITS      , "radians m-1"  ))
    call check(nf90_put_att(ncid, lat_varid       , UNITS      , "degrees_north"))
    call check(nf90_put_att(ncid, lon_varid       , UNITS      , "degrees_east" ))

    call check(nf90_def_var(ncid, "div"           , nf90_double, (/lon_dimid, lat_dimid/), div_varid))
    call check(nf90_def_var(ncid, "vor"           , nf90_double, (/lon_dimid, lat_dimid/), vor_varid))
    call check(nf90_def_var(ncid, "power_div"     , nf90_double, (/wavenumber_dimid/), powerdiv_varid))
    call check(nf90_put_att(ncid, powerdiv_varid  , UNITS      , "J kg-1"))
    call check(nf90_def_var(ncid, "power_vor"     , nf90_double, (/wavenumber_dimid/), powervor_varid))
    call check(nf90_put_att(ncid, powervor_varid  , UNITS      , "J kg-1"))
    call check(nf90_def_var(ncid, "wavelength"    , nf90_double, (/wavenumber_dimid/), wavelength_varid))
    call check(nf90_put_att(ncid, wavelength_varid, UNITS      , "km"    ))
    call check(nf90_enddef(ncid))
   
    call check(nf90_put_var(ncid, wavenumber_varid, wavenumber))
    call check(nf90_put_var(ncid, lat_varid       , lats(2:nlats_in-1)))
    call check(nf90_put_var(ncid, lon_varid       , lons(2:nlons_in-1)))
    call check(nf90_put_var(ncid, div_varid       , div      ))
    call check(nf90_put_var(ncid, vor_varid       , vor      ))
    call check(nf90_put_var(ncid, powerdiv_varid  , power_div))
    call check(nf90_put_var(ncid, powervor_varid  , power_vor))
    call check(nf90_put_var(ncid, wavelength_varid, wavelength))
    call check(nf90_close(ncid))

  end subroutine output

  subroutine calc_div_vor(u, v, div, vor)
  
    real(r8), intent(in) :: u(0:nlons+1,0:nlats+1), v(0:nlons+1,0:nlats+1)
    real(r8), intent(inout) :: div(1:nlons,1:nlats), vor(1:nlons,1:nlats)
    integer i, j
    
    do j = 1, nlats
      do i = 1, nlons
        div(i,j) = 1.0 / radius / cos(rad * lats(j+1)) * ((u(i+1,j) - u(i-1,j)) / (2 * dlon) + &
                  (v(i,j+1) * cos(rad * lats(j+2)) - &
                   v(i,j-1) * cos(rad * lats(j  ))) / (2 * dlat))  ! lats's index needs + 1 corresponding to div and vor 

        vor(i,j) = 1.0 / radius / cos(rad * lats(j+1)) * ((v(i+1,j) - v(i-1,j)) / (2 * dlon) - &
                  (u(i,j+1) * cos(rad * lats(j+2)) - &
                   u(i,j-1) * cos(rad * lats(j  ))) / (2 * dlat))
      end do
    end do

  end subroutine calc_div_vor
  
  subroutine calc_DCT(div, vor, DCT_div, DCT_vor)

    real(r8), intent(in) :: div(1:nlons,1:nlats), vor(1:nlons,1:nlats)
    real(r8), intent(inout) :: DCT_div(0:nlons-1,0:nlats-1), &
                               DCT_vor(0:nlons-1,0:nlats-1)
    real(r8), allocatable :: betam(:), betan(:) 
    real(r8) tmp
    integer i, j, m, n
    
    allocate(betam(0:nlons-1),&
             betan(0:nlats-1))

    betam(0) = sqrt(1.0 / nlons)
    betam(1:nlons-1) = sqrt(2.0 / nlons)
    betan(0) = sqrt(1.0 / nlats)
    betan(1:nlats-1) = sqrt(2.0 / nlats)
    
    do n = 0, nlats - 1
      do m = 0, nlons - 1
        DCT_div(m,n) = 0.0
        DCT_vor(m,n) = 0.0
        do j = 0, nlats - 1
          do i = 0, nlons - 1
            tmp = cos(pi * m * (i + 0.5) / nlons) * &
                  cos(pi * n * (j + 0.5) / nlats)
            DCT_div(m,n) = DCT_div(m,n) + div(i+1,j+1) * tmp 
            DCT_vor(m,n) = DCT_vor(m,n) + vor(i+1,j+1) * tmp
          end do
        end do
        DCT_div(m,n) = betam(m) * betan(n) * DCT_div(m,n)
        DCT_vor(m,n) = betam(m) * betan(n) * DCT_vor(m,n)
      end do
    end do
    deallocate(betam, betan)

  end subroutine calc_DCT

  subroutine calc_power_spectral(dct_div, dct_vor, power_div, power_vor, wavelength, wavenumber)
    
    real(r8), intent(in) :: dct_div(0:nlons-1, 0:nlats-1), &
                            dct_vor(0:nlons-1, 0:nlats-1)
    real(r8), intent(inout) :: power_div(Nwave), power_vor(Nwave), &
                              wavelength(Nwave), wavenumber(Nwave)
    real(r8) mu_lower, mu_upper, mu, khat, nlonlat
    integer k, m, n
    
    nlonlat = nlons * nlats

    do k = 1, Nwave
      power_div(k) = 0.0
      power_vor(k) = 0.0
      mu_lower = real(k) / min(nlons, nlats)
      mu_upper = real(k + 1) / min(nlons, nlats)
      wavelength(k) = 2.0 * (radius * dlon * 1.e-03) * min(nlons, nlats) / k !km
      wavenumber(k) = pi * mu_lower / (radius * dlon)
      do n = 0, nlats - 1
        do m = 0, nlons - 1
          if (m == 0 .and. n == 0) cycle
          mu = sqrt(real(m)**2 / real(nlons)**2 + real(n)**2 / real(nlats)**2)
          khat = pi * mu / (radius * dlon)
          if (mu >= mu_lower .and. mu < mu_upper) then
            ! based on divergence and vorticity
            if (vor_div_based) then
              power_div(k) = power_div(k) + 0.5 * dct_div(m,n)**2 / nlonlat / khat**2
              power_vor(k) = power_vor(k) + 0.5 * dct_vor(m,n)**2 / nlonlat / khat**2
            else
            ! based on u and v
              power_div(k) = power_div(k) + 0.5 * dct_div(m,n)**2 / nlonlat * (radius * dlon) * min(nlons, nlats) / pi
              power_vor(k) = power_vor(k) + 0.5 * dct_vor(m,n)**2 / nlonlat * (radius * dlon) * min(nlons, nlats) / pi
            end if
          end if
        end do
      end do
    end do

  end subroutine calc_power_spectral

  subroutine check(status)
    
    integer, intent(in):: status

    if (status /= nf90_noerr) then
      print*, trim(nf90_strerror(status))
      stop "stoped"
    end if

  end subroutine check

end program calc_kinetic_spectra
